package com.lglbc.aidemo;

import jakarta.annotation.PostConstruct;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.i18n.LocaleContextHolder;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.io.PathResource;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.List;
import java.util.Map;

@RestController
public class AiController {

    @Autowired
    private ToolCallbackProvider toolCallbackProvider;
    @Autowired
    private ChatModel zhiPuAiChatModel;
    @Autowired
    private ChatClient zpChatGlm4vFlashClient;
    @Autowired
    private ChatClient zpChatGlm4PlusClient;
    @Autowired
    private VectorStore vectorStore;
    @Autowired
    private ChatMemory chatMemory;
    @Autowired
    private DocumentTransformer documentTransformer;

    @GetMapping("/ai")
    public String generation(@RequestParam(defaultValue = "你觉得知识能改变命运么？") String userInput) {
        ChatClient chatClient = ChatClient.create(zhiPuAiChatModel);
        return chatClient.prompt(userInput).call().content();
    }

    @GetMapping("/ai/json")
    public Object generationJson(@RequestParam(defaultValue = "随机生成一份用户信息") String userInput) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        User zpContent = zpChatClient.prompt(userInput)
                .call().entity(User.class);
        return zpContent;
    }

    @GetMapping("/ai/json/list")
    public Object generationJsonList(@RequestParam(defaultValue = "随机生成5个用户信息") String userInput) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        List<User> zpContent = zpChatClient.prompt(userInput)
                .call().entity(new ParameterizedTypeReference<>() {
                });
        return zpContent;
    }

    @GetMapping(value = "/ai/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE + ";charset=UTF-8")
    public Flux<String> streamContent(@RequestParam(defaultValue = "和我说一个长笑话") String userInput) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        return zpChatClient.prompt(userInput)
                .advisors(new SimpleLoggerAdvisor())
                .stream().content();
    }

    @GetMapping("/ai/prompt/template")
    public Object promptTemplate(@RequestParam(defaultValue = "乐哥聊编程") String username) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        return zpChatClient.prompt()
                .user(u -> u.text("请帮我写一首诗，作者是{userName},内容围绕作者的名字去写")
                        .param("userName", username)
                )
                .advisors(new SimpleLoggerAdvisor())
                .call().content();
    }

    @GetMapping("/ai/prompt/template/map")
    public Object promptTemplateMap(@RequestParam(defaultValue = "乐哥聊编程") String username) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        return zpChatClient.prompt()
                .user(u -> u.text("请帮我写一首诗，作者是{userName},内容围绕作者的名字去写")
                        .params(Map.of("userName", username)))
                .advisors(new SimpleLoggerAdvisor())
                .call().content();
    }

    @GetMapping("/ai/client/configuration/model1")
    public Object clientConfiguration1() {
        // Create ChatClient instances programmatically
        return zpChatGlm4PlusClient.prompt("你使用了哪种大模型？")
                .call().content();
    }

    @GetMapping("/ai/client/configuration/model2")
    public Object clientConfiguration2() {
        // Create ChatClient instances programmatically
        return zpChatGlm4vFlashClient.prompt("你使用了哪种大模型？")
                .call().content();
    }

    @GetMapping("/ai/chat/memory")
    public Object chatMemory(String userInput) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        return zpChatClient.prompt(userInput)
                .advisors(new SimpleLoggerAdvisor(), MessageChatMemoryAdvisor.builder(chatMemory).build())
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, "conversationId"))
                .call().content();
    }

    @GetMapping("/ai/tools")
    public Object tools(String userInput) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        return zpChatClient.prompt(userInput)
                .advisors(new SimpleLoggerAdvisor())
                .tools(new DateTimeTools())
                .call().content();
    }


    @GetMapping("/ai/rag")
    public Object rag(@RequestParam(defaultValue = "ElasticSearch 如何实现聚合？") String userInput) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        return zpChatClient.prompt(userInput)
                .advisors(new SimpleLoggerAdvisor(), new QuestionAnswerAdvisor(vectorStore))
                .call().content();
    }


    @GetMapping("/ai/mcp")
    public Object mcp(String userInput) {
        // Create ChatClient instances programmatically
        ChatClient zpChatClient = ChatClient.create(zhiPuAiChatModel);
        return zpChatClient.prompt(userInput)
                .advisors(new SimpleLoggerAdvisor())
                .toolCallbacks(toolCallbackProvider)
                .tools(new DateTimeTools())
                .call().content();
    }



    record User(String name, String age, String sex) {

    }

    static class DateTimeTools {

        @Tool(description = "获取当前时间")
        String getCurrentDateTime() {
            return LocalDateTime.now().atZone(LocaleContextHolder.getTimeZone().toZoneId()).toString();
        }

        @Tool(description = "获取用户的生日，入参是用户名,出参是年月日 如 1994年1月1日")
        String getBirthdayDate(String userName) {
            if (userName.equals("张三")) {
                return "1994年8月29日";
            } else if (userName.equals("李斯")) {
                return "1994年4月29日";
            } else {
                return "没有这个人的信息";
            }
        }
    }
    @PostConstruct
    public void init() {
        Resource resource = new PathResource("/Users/ly/Desktop/MySQL.pdf");
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(resource);
        List<Document> documents = tikaDocumentReader.get();
        List<Document> transformedDocs = documentTransformer.apply(documents);
        vectorStore.accept(transformedDocs);
        System.out.println("数据初始化完成");
    }

}
